import torch

# test_dtype = torch.bfloat16
# test_device = "cuda"
test_dtype = torch.float32
test_device = "cpu"

from sgm.modules.diffusionmodules.dit import DiffusionTransformer
from huggingface_hub import hf_hub_download
from diffusers import CogView3PlusTransformer2DModel
from yiyi_convert_cogview3_to_diffusers import convert_cogview3_transformer_checkpoint_to_diffusers

# (1). create common inputs for testing

x = torch.load("dit_x.pt").to(test_device).to(test_dtype) # [2, 16, 128, 128]
timesteps = torch.load("dit_timesteps.pt").to(test_device).to(test_dtype) # [2]
context = torch.load("dit_context.pt").to(test_device).to(test_dtype) #[2, 224, 4096]
y = torch.load("dit_y.pt").to(test_device).to(test_dtype) #[2, 1536]
kwargs = {'target_size': [(1024, 1024)], "idx": timesteps}

# (2). load and clean up state_dict
repo_id = "ZP2HF/CogView3-SAT"
filename ="cogview3plus_3b/1/mp_rank_00_model_states.pt"
ckpt_path_cogview3_plus = hf_hub_download(repo_id=repo_id, filename=filename)
state_dict = torch.load(ckpt_path_cogview3_plus)["module"]
state_dict = {k.replace("model.diffusion_model.", ""): v for k, v in state_dict.items()}

# (3) make model(original implementation)
model_config = {
    "in_channels": 16,
    "out_channels": 16,
    "hidden_size": 2560,
    "patch_size": 2,
    "num_layers": 30,
    "num_attention_heads": 64,
    "text_length": 224,
    "time_embed_dim": 512,
    "num_classes": "sequential",
    "adm_in_channels": 1536,
    "modules": {
        'pos_embed_config': {
            'target': 'sgm.modules.diffusionmodules.dit.PositionEmbeddingMixin',
            'params': {'max_height': 128, 'max_width': 128, 'max_length': 4096}
        },
        'patch_embed_config': {
            'target': 'sgm.modules.diffusionmodules.dit.ImagePatchEmbeddingMixin',
            'params': {'text_hidden_size': 4096}
        },
        'attention_config': {
            'target': 'sgm.modules.diffusionmodules.dit.AdalnAttentionMixin',
            'params': {'qk_ln': True}
        },
        'final_layer_config': {
            'target': 'sgm.modules.diffusionmodules.dit.FinalLayerMixin'
        }
    },
    "layernorm_order": "pre",
    "activation_func": None,
    "elementwise_affine": False,
    "parallel_output": True,
    "block_size": 16,
    "dtype": "fp32" if test_dtype == torch.float32 else "bf16"
}
model_sat = DiffusionTransformer(**model_config)
model_sat.to(test_device).to(test_dtype)
missing_keys, unexpected_keys = model_sat.load_state_dict(state_dict, strict=False)
print(f"missing_keys: {missing_keys}")
print(f"unexpected_keys: {unexpected_keys}")

# run a forward pass
model_sat.eval()
with torch.no_grad():
    out = model_sat(
        x=x, 
        timesteps=timesteps,
        context=context, 
        y=y, 
        **kwargs)

# (4) make model(diffusers)
converted_state_dict = convert_cogview3_transformer_checkpoint_to_diffusers(state_dict)

model = CogView3PlusTransformer2DModel()
model.to(test_device).to(test_dtype)
model.load_state_dict(converted_state_dict, strict=True)

model.eval()
with torch.no_grad():
    out_d = model(
        hidden_states=x, 
        timestep=timesteps, 
        encoder_hidden_states=context, 
        pooled_projections=y, 
        return_dict=False)[0]

print(f" output shape: {out_d.shape}")
print(f" expected output shape: {out.shape}")
assert out_d.shape == out.shape
assert (out_d - out).abs().max() < 1e-4, f"max diff: {(out_d - out).abs().max()}"